// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/benchmarks/ExpressionBenchmarkBuilder.h>
#include <melon/benchmark.h>
#include <pollux/testing/vector/vector_test_base.h>

namespace kumo::pollux {

ExpressionBenchmarkSet& ExpressionBenchmarkSet::addExpression(
    const std::string& name,
    const std::string& expression) {
  expressions_.push_back(
      std::make_pair(name, builder_.compileExpression(expression, inputType_)));
  return *this;
}

ExpressionBenchmarkSet& ExpressionBenchmarkSet::addExpressions(
    const std::vector<std::pair<std::string, std::string>>& expressions) {
  for (auto& [name, expression] : expressions) {
    addExpression(name, expression);
  }
  return *this;
}

// Make sure all input vectors are generated.
void ExpressionBenchmarkBuilder::ensureInputVectors() {
  for (auto& [_, benchmarkSet] : benchmarkSets_) {
    if (!benchmarkSet.inputRowVector_) {
      VectorFuzzer fuzzer(benchmarkSet.fuzzerOptions_, pool());
      benchmarkSet.inputRowVector_ = std::dynamic_pointer_cast<RowVector>(
          fuzzer.fuzzFlat(benchmarkSet.inputType_));
    }
  }
}

void ExpressionBenchmarkBuilder::testBenchmarks() {
  ensureInputVectors();

  auto evalExpression = [&](auto& exprSet, auto& inputVector) {
    exec::EvalCtx evalCtx(&this->execCtx_, &exprSet, inputVector.get());
    SelectivityVector rows(inputVector->size());
    std::vector<VectorPtr> results(1);
    exprSet.eval(rows, evalCtx, results);
    return results[0];
  };

  auto testBenchmarkSet = [&](auto& benchmarkSet) {
    if (benchmarkSet.expressions_.size() == 0) {
      return;
    }
    // Evaluate the first expression.
    auto it = benchmarkSet.expressions_.begin();
    auto refResult = evalExpression(it->second, benchmarkSet.inputRowVector_);
    it++;
    while (it != benchmarkSet.expressions_.end()) {
      auto result = evalExpression(it->second, benchmarkSet.inputRowVector_);
      test::assertEqualVectors(refResult, result);
      it++;
    }
  };

  for (auto& [_, benchmarkSet] : benchmarkSets_) {
    if (!benchmarkSet.disableTesting_) {
      testBenchmarkSet(benchmarkSet);
    }
  }
}

void ExpressionBenchmarkBuilder::registerBenchmarks() {
  ensureInputVectors();
  // Generate input vectors if needed.
  for (auto& [setName, benchmarkSet] : benchmarkSets_) {
    for (auto& [exprName, exprSet] : benchmarkSet.expressions_) {
      auto name = fmt::format("{}##{}", setName, exprName);
      auto& inputVector = benchmarkSet.inputRowVector_;
      auto times = benchmarkSet.iterations_;
      // The compiler does not allow capturing exprSet int the lambda.
      auto& exprSetLocal = exprSet;
      melon::addBenchmark(
          __FILE__, name, [this, &inputVector, &exprSetLocal, times]() {
            int cnt = 0;
            melon::BenchmarkSuspender suspender;
            // TODO: shall we cache those.
            exec::EvalCtx evalCtx(
                &this->execCtx_, &exprSetLocal, inputVector.get());
            SelectivityVector rows(inputVector->size());
            suspender.dismiss();

            std::vector<VectorPtr> results(1);
            for (auto i = 0; i < times; i++) {
              exprSetLocal.eval(rows, evalCtx, results);

              // TODO: add flag to enable/disable flattening.
              BaseVector::flattenVector(results[0]);

              // TODO: add flag to enable/disable reuse.
              results[0]->prepare_for_reuse();

              cnt += results[0]->size();
            }
            melon::doNotOptimizeAway(cnt);
            return 1;
          });
    }

    melon::addBenchmark(__FILE__, "-", []() -> unsigned { return 0; });
  }
}
} // namespace kumo::pollux
